#!/usr/bin/python

"""
Usage: $shell precondition.py [OPTIONS]
              OPTIONS:
                --patchall, -p : preparation for patchall
                --upgrade,  -u : preparation for upgrade
                --porting,  -t : preparation for porting
"""

__author__ = 'duanqz@gmail.com'



import shutil
import subprocess
import os, sys
import commands
import getopt
from config import Config
from changelist import ChangeList
from formatters.log import Log, Paint
from reverses.zipformatter import ZipFormatter



TAG="precondition"

FRAMEWORK_JARS = ("framework.jar", "services.jar", "telephony-common.jar", "wifi-service.jar", "android.policy.jar",
                  "secondary-framework.jar", "secondary_framework.jar",
                  "framework-ext.jar", "framework_ext.jar",
                  "framework2.jar",
                  "mediatek-framework.jar")

PARTITIONS     = ("secondary_framework.jar.out", "secondary-framework.jar.out",
                  "framework-ext.jar.out", "framework_ext.jar.out",
                  "framework2.jar.out")


USELESS_DIRS   = ("framework.jar.out/smali/com/flyme", "framework.jar.out/smali/com/meizu",
                  "framework.jar.out/smali/flyme", "framework.jar.out/smali/meizu",
                  "framework.jar.out/original",
                  "services.jar.out/original",
                  "telephony-common.jar.out/original",
                  "wifi-service.jar.out/original",
                  "android.policy.jar.out/original",
                  "framework-res/original")


class Options(object):

    def __init__(self):
        self.prepare = True                    # Whether to trigger prepare action
        self.patchXml = Config.PATCHALL_XML    # The patch XML, default to be PATCHALL_XML
        self.baseName = "base"                 # Short base device name, default to be 'base'
        self.commit1  = None                   # The 7 bits lower commit ID
        self.commit2  = None                   # The 7 bits upper commit ID
        self.olderRoot = None                  # The older and newer is a pair of directory for comparing
        self.newerRoot = None                  # Default to be AOSP and BOSP

    @staticmethod
    def usage():
        print __doc__
        sys.exit()

    def handle(self, argv):
        if len(argv) == 1: Options.usage()

        try:
            (opts, args) = getopt.getopt(argv[1:], "hlputb:c1:c2:", \
                            [ "help", "loosely", "patchall", "upgrade", "porting", "base=", "commit1=", "commit2=" ])
            Log.d(TAG, "Program args = %s" %args)
        except getopt.GetoptError:
            Options.usage()

        for name, value in opts:
            if name in ("--help", "-h"):
                Options.usage()

            elif name in ("--loosely", "-l"):
                self.prepare = False

            elif name in ("--patchall", "-p"):
                self.patchXml  = Config.PATCHALL_XML
                self.olderRoot = Config.AOSP_ROOT
                self.newerRoot = Config.BOSP_ROOT

            elif name in ("--upgrade", "-u"):
                self.patchXml  = Config.UPGRADE_XML
                self.olderRoot = Config.LAST_BOSP_ROOT
                self.newerRoot = Config.BOSP_ROOT

            elif name in ("--porting", "-t"):
                # The older and newer root are generated by the commit1 and commit2
                self.patchXml = Config.PORTING_XML

            elif name in ("--base", "-b"):
                if len(value) > 0: self.baseName = value

            elif name in ("--commit1", "-1"):
                if len(value) > 0: self.commit1 = value

            elif name in ("--commit2", "-2"):
                if len(value) > 0: self.commit2 = value

        self.dump()

        return self


    def dump(self):
        Log.d(TAG, "prepare = %s" %str(self.prepare))
        Log.d(TAG, "baseName = %s, patchXml = %s, olderRoot = %s, newerRoot = %s, commit1 = %s, commit2 = %s"
              %(self.baseName, self.patchXml, self.olderRoot, self.newerRoot, self.commit1, self.commit2))


# Global OPTIONS
OPTIONS = Options()

class Prepare:

    def __init__(self):
        """ baseName is the short device name
        """

        Log.i(TAG, "Start preparing essential files in %s" %Config.AUTOPATCH)

        self.baseDevice = BaseDevice(OPTIONS.baseName)
        if   OPTIONS.patchXml == Config.PATCHALL_XML: self.patchall()
        elif OPTIONS.patchXml == Config.UPGRADE_XML:  self.upgrade()
        elif OPTIONS.patchXml == Config.PORTING_XML:  self.porting()


    def patchall(self):
        """ Prepare precondition of patchall
        """

        # Phase 1: prepare AOSP
        self.baseDevice.aosp(OPTIONS.olderRoot)

        # Phase 2: prepare BOSP
        self.baseDevice.bosp(OPTIONS.newerRoot)

        # Phase 3: record last head
        self.baseDevice.setLastHead()

        # Phase 4: prepare patch XML, not force
        ChangeList(OPTIONS.olderRoot, OPTIONS.newerRoot, OPTIONS.patchXml).make(force=False)


    def upgrade(self):
        """ Prepare precondition of upgrade
        """

        # Remove last_bosp and bosp
        Utils.run(["rm", "-rf", OPTIONS.olderRoot], stdout=subprocess.PIPE).communicate()
        Utils.run(["rm", "-rf", OPTIONS.newerRoot], stdout=subprocess.PIPE).communicate()

        lastBoardZip  = os.path.join(Config.PRJ_ROOT, "board/last_board.zip")
        boardZip      = os.path.join(Config.PRJ_ROOT, "board/board.zip")

        if os.path.exists(lastBoardZip) and os.path.exists(boardZip):
            # Phase 1: prepare LAST_BOSP from last_board.zip
            Utils.decode(lastBoardZip, OPTIONS.olderRoot)

            # Phase 2: prepare BOSP from board.zip
            Utils.decode(boardZip, OPTIONS.newerRoot)

            # Phase 3: prepare patch XML
            ChangeList(OPTIONS.olderRoot, OPTIONS.newerRoot, OPTIONS.patchXml).make(force=True)

        else:
            if OPTIONS.commit1 is not None:
                self.baseDevice.setLastHead(OPTIONS.commit1)

            # Phase 1: get last and origin head from base device
            (lastHead, origHead) = self.baseDevice.getLastAndOrigHead()
            if lastHead == origHead:
                Log.w(TAG, Paint.red("Nothing to upgrade. Did you forget to sync the %s ?" % OPTIONS.baseName))
                Log.w(TAG, Paint.red("In the root directory, try the following command to sync:"))
                Log.w(TAG, Paint.red("  $ repo sync devices/%s" % OPTIONS.baseName))

            # Phase 2: porting from last to origin head
            OPTIONS.commit1 = lastHead[0:7]
            OPTIONS.commit2 = origHead[0:7]
            self.porting()


    def porting(self):
        """ Porting changes from commit1 to commit2
            commit is 7 bits
        """

        # Phase 1: get the lower and upper commit
        commitModel = CommitModel(self.baseDevice)
        (lowerCommit, upperCommit) = commitModel.getCommitRange(OPTIONS.commit1, OPTIONS.commit2)
        Log.d(TAG, "Porting.prepare(). lowerCommit = %s, upperCommit = %s" %(lowerCommit, upperCommit))

        # Phase 2: Prepare the older and newer root
        if OPTIONS.patchXml == Config.PORTING_XML:
            OPTIONS.newerRoot = os.path.join(Config.AUTOPATCH, "%s_newer_%s" % (self.baseDevice.name(), upperCommit))
            OPTIONS.olderRoot = os.path.join(Config.AUTOPATCH, "%s_older_%s" % (self.baseDevice.name(), lowerCommit))
        filesChanged = self.baseDevice.getFilesChanged(lowerCommit, upperCommit)
        self.baseDevice.portingFiles(upperCommit, filesChanged, OPTIONS.newerRoot)
        self.baseDevice.portingFiles(lowerCommit, filesChanged, OPTIONS.olderRoot)

        # Phase 3: Restore the commit model
        commitModel.restore()

        # Phase 4: prepare patch XML
        ChangeList(OPTIONS.olderRoot, OPTIONS.newerRoot, OPTIONS.patchXml).make(force=True)



class CommitModel:
    """ The Commits Model of a device
    """

    def __init__(self, baseDevice):
        """ Initialize the commits model from a device path.
        """

        self.baseDevice = baseDevice
        self.commitIDs = []
        self.comments  = []

        baseDevice.fillCommitsInfo(self.commitIDs, self.comments)


    def showUserInputHint(self):
        """ Show user input hint
        """

        for i in range(0, len(self.commitIDs)):
            commitID = self.commitIDs[i]
            comment  = self.comments[i]
            print "  %s %s" % (Paint.bold(commitID), comment)

        deviceName = self.baseDevice.name()
        oneCommit  = Paint.bold(self.commitIDs[0])
        twoCommits = "%s %s" % (Paint.bold(self.commitIDs[-1]), Paint.bold(self.commitIDs[0]))
        print "  ________________________________________________________________________________________"
        print "                                                                                          "
        print "  Each 7 bits SHA1 code identify a commit on %s," % Paint.blue(deviceName),
        print "  You could input:                                                                        "
        print "  - Only one single commit, like: %s" % oneCommit
        print "    will porting changes between the selected and the latest from %s to your device" % Paint.blue(deviceName) 
        print "  - Two commits as a range, like: %s" % twoCommits
        print "    will porting changes between the two selected from %s to your device" % Paint.blue(deviceName)
        print "  ________________________________________________________________________________________"
        print "                                                                                          "


    def readUserInput(self):
        """ Read user input
        """

        self.showUserInputHint()

        userInput = raw_input(Paint.bold(">>> Input the 7 bits SHA1 commit ID (q to exit):  "))
        if userInput in ("q", "Q"): sys.exit()

        commits = userInput.split()
        size = len(commits)
        commit1 = commit2 = None
        if size > 0 : commit1 = commits[0]
        if size > 1 : commit2 = commits[1]

        return (commit1, commit2)


    def computeLowerAndUpper(self, commit1, commit2=None):
        """ Retrieve the lower and upper commit ID
        """

        if commit2 == None: commit2 = commit1

        try:    index1 = self.commitIDs.index(commit1)
        except: index1 = 0

        try:    index2 = self.commitIDs.index(commit2)
        except: index2 = 0

        if  index1 == index2:
            upper = 0
            lower = index1
        elif index1 < index2:
            upper = index1
            lower = index2
        else:
            upper = index2
            lower = index1

        lowerCommit = self.commitIDs[lower]
        upperCommit = self.commitIDs[upper]

        return (lowerCommit, upperCommit)


    def getCommitRange(self, commit1, commit2):
        """ Get the range of commit1 and commit2
        """

        # If no commit present, ask for user input
        if commit1 == None and commit2 == None:
            (commit1, commit2) = self.readUserInput()

        return self.computeLowerAndUpper(commit1, commit2)


    def restore(self):
        """ Restore the base device
        """

        self.baseDevice.reset(self.commitIDs[0])

# End of class PortingDevice

class BaseDevice:

    def __init__(self, baseName):
        self.baseName = baseName
        self.basePath = BaseDevice.getBasePath(baseName)
        self.targetPath = os.path.abspath(Config.PRJ_ROOT)

        targetName = os.path.basename(self.targetPath)
        self.lastHeadPath = os.path.join(self.basePath, ".git/%s:LAST_HEAD" %targetName)
        self.origHeadPath = os.path.join(self.basePath, ".git/%s:ORIG_HEAD" %targetName)

    def name(self):
        return self.baseName

    @staticmethod
    def getBasePath(base):
        try:
            devices = os.path.join(os.environ["PORT_ROOT"], "devices")
            return os.path.join(devices, base)
        except KeyError:
            Log.e(TAG, "device %s not found" % base)
            sys.exit(155)


    def setLastHead(self, commit=None):
        """ Set last head of input device
        """

        if commit == None: commit = self.parseHeadCommit()
        BaseDevice.writeCommit(self.lastHeadPath, commit)


    def getLastAndOrigHead(self):
        """ Get the last head and the origin head of input device.
        """

        # If last or orig head not exists, write the current head into them.
        head = self.parseHeadCommit()
        if not os.path.exists(self.lastHeadPath): BaseDevice.writeCommit(self.lastHeadPath, head)
        if not os.path.exists(self.origHeadPath): BaseDevice.writeCommit(self.origHeadPath, head)

        # Check whether need to update the lastHead and origHead
        oldOrigHead = BaseDevice.readCommit(self.origHeadPath)
        newOrigHead = head

        if oldOrigHead == newOrigHead:
            Log.d(TAG, "BaseDevice.getLastAndOrigHead(). oldOrig == newOrig")
            pass
        else:
            Log.d(TAG, "BaseDevice.getLastAndOrigHead(). oldOrig -> LAST_HEAD, newOrig -> ORIG_HEAD")
            BaseDevice.writeCommit(self.lastHeadPath, oldOrigHead)
            BaseDevice.writeCommit(self.origHeadPath, newOrigHead)

        lastHead = BaseDevice.readCommit(self.lastHeadPath)
        origHead = BaseDevice.readCommit(self.origHeadPath)
        Log.d(TAG, "BaseDevice.getLastAndOrigHead(). lastHead = %s, origHead = %s" %(lastHead, origHead))

        return (lastHead, origHead)


    def parseHeadCommit(self):
        """ Parse out current head commit
        """

        os.chdir(self.basePath)

        subp = Utils.run(["git", "rev-parse", "HEAD"], stdout=subprocess.PIPE)

        os.chdir(self.targetPath)

        return subp.stdout.read().strip("\n")


    def fillCommitsInfo(self, commitIDs, comments):
        """ Fill the commitIDs and comments
            commitIDs is an array of 7 bits SHA1 code.
            comment is an array of string.
        """

        os.chdir(self.basePath)

        subp = Utils.run(["git", "log", "--oneline"], stdout=subprocess.PIPE)
        while True:
            buff = subp.stdout.readline().strip('\n')
            if buff == '' and subp.poll() != None:
                break

            buff = buff.strip()
            # The first 7 bits is commit ID
            commitIDs.append(buff[0:7])
            comments.append(buff[7:])

        os.chdir(self.targetPath)


    def reset(self, commit):
        """ Reset to commit
        """

        os.chdir(self.basePath)

        subp = Utils.run(["git", "reset", "--hard", commit], stdout=subprocess.PIPE)
        subp.communicate()

        os.chdir(self.targetPath)


    def aosp(self, aospDst):
        """ Prepare AOSP to asopDst
        """

        aospSrc = os.path.join(self.basePath, "vendor/aosp")
        # If no AOSP under vendor/ , decode them out
        if not os.path.exists(aospSrc):
            os.makedirs(aospSrc)
            vendorRoot = os.path.join(self.basePath, "vendor")
            Utils.decodeAPKandJAR(vendorRoot, aospSrc)

        if not os.path.exists(aospDst):
            Log.i(TAG, "Generating %s from %s" % (aospDst, aospSrc))
            Utils.copyAPKandJAR(aospSrc, aospDst)


    def bosp(self, bospDst, force=True):
        """ Prepare BOSP, set force to be False to not generate again if exists.
        """

        if force:
            subp = Utils.run(["rm", "-rf", bospDst], stdout=subprocess.PIPE)
            subp.communicate()

        Log.i(TAG, "Generating %s from %s" %(bospDst, self.basePath))
        Utils.copyAPKandJAR(self.basePath, bospDst)


    def getFilesChanged(self, lowerCommit, upperCommit):
        """ Get files changed from lower to upper commit
        """

        os.chdir(self.basePath)

        changes = []
        subp = Utils.run(["git", "diff", "--name-only", lowerCommit, upperCommit], stdout=subprocess.PIPE)
        while True:
            buff = subp.stdout.readline().strip('\n')
            if buff == '' and subp.poll() != None:
                break

            changes.append(buff.strip())

        os.chdir(self.targetPath)

        return changes


    def portingFiles(self, commit, filesChanged, dstDir):
        """ Generate patch files for porting
        """

        # Reset to the commit
        self.reset(commit)

        Log.i(TAG, "Generating %s from %s at commit %s" % (dstDir, self.basePath, commit))
        if not os.path.exists(dstDir): os.makedirs(dstDir)

        # Copy changed items from source
        for item in filesChanged:
            src = os.path.join(self.basePath, item)
            dst = os.path.join(dstDir, item)
            if os.path.exists(src):

                # Only copy files in FRAMEWORK_JARS
                if not Utils.isInFramework(item): continue

                # We need to format the SMALI file.
                # Copy all the sub classes even if no change.
                Utils.copyWholly(src, os.path.dirname(dst))

        Utils.combineFrameworkPartitions(dstDir)


    @staticmethod
    def readCommit(filepath):
        handle = open(filepath, "r")
        content = handle.read().strip("\n")
        handle.close()

        return content

    @staticmethod
    def writeCommit(filepath, commit):
        handle = open(filepath, "w")
        handle.write(commit)
        handle.close()


class Utils:
    """ Utilities
    """

    @staticmethod
    def decode(boardZip, out):
        """ Decode FRAMEWORK_JARS in board.zip into out directory.
        """

        Log.i(TAG, "Generating %s from %s" %(out, boardZip))

        # Phase 1: normalize
        temp = Utils.otaNormalize(boardZip)
        if temp == None:
            Log.e(TAG, "decode(): ota normalized failed")
            return

        if not os.path.exists(out): os.makedirs(out)

        Utils.decodeAPKandJAR(temp, out)

        shutil.rmtree(temp)

        # Phase 3: combine framework partitions
        Utils.combineFrameworkPartitions(out)

        # Phase 4: remove useless files
        Utils.removeUseless(out)


    @staticmethod
    def decodeAPKandJAR(root, out):
        # Format path
        if os.path.exists(os.path.join(root, "SYSTEM")):
            shutil.move(os.path.join(root, "SYSTEM"), os.path.join(root, "system"))

        dirname = os.path.join(root, "system/framework")

        Log.i(TAG, "decoding framework-res.apk")
        jarpath = os.path.join(dirname, "framework-res.apk")
        jarout  = os.path.join(out, "framework-res")
        subp = Utils.run(["apktool", "d", "-f", jarpath, "-o", jarout], stdout=subprocess.PIPE)
        Utils.printSubprocessOut(subp)

        for jarname in FRAMEWORK_JARS:
            jarpath = os.path.join(dirname, jarname)
            if os.path.exists(jarpath):
                Log.i(TAG, "decoding %s" % jarname)
                jarout = os.path.join(out, jarname + ".out")
                subp = Utils.run(["apktool", "d", "-f", jarpath, "-o", jarout], stdout=subprocess.PIPE)
                Utils.printSubprocessOut(subp)


    @staticmethod
    def otaNormalize(boardZip):
        """ Normalize the OTA package zip, return the root directory of unziped files
        """

        if not os.path.exists(boardZip):
            Log.e(TAG, "oatNormalize() % not exists" % boardZip)
            return None

        zipFormatter = ZipFormatter.create(ZipFormatter.genOptions(boardZip))
        # Do not need to zip back
        zipFormatter.format(zipBack=False)

        root = zipFormatter.getFilesRoot()
        if not os.path.exists(root):
            Log.e(TAG, "otaNormalize() normalize %s failed!" % boardZip)
            return None

        return root


    @staticmethod
    def copyAPKandJAR(src, dst):
        if not os.path.exists(dst):
            os.makedirs(dst)

        bootImage = os.path.join(src, "boot.img.out")
        subp = Utils.run(["cp", "-r", bootImage, dst], stdout=subprocess.PIPE)
        subp.communicate()

        frwRes = os.path.join(src, "framework-res")
        subp = Utils.run(["cp", "-r", frwRes, dst], stdout=subprocess.PIPE)
        subp.communicate()

        for jarname in FRAMEWORK_JARS:
            jarname += ".out"
            srcJar = os.path.join(src, jarname)
            if os.path.exists(srcJar):
                Log.d(TAG, "Utils.copyAPKandJAR(). copying %s to %s" %(srcJar, dst))
                subp = Utils.run(["cp", "-r", srcJar, dst], stdout=subprocess.PIPE)
                subp.communicate()

        Utils.combineFrameworkPartitions(dst)


    @staticmethod
    def combineFrameworkPartitions(frameworkDir):
        """ Combine framework partitions into framework.jar.out.
        """

        # For Android 5.0, handle dex partitions.
        dst = os.path.join(frameworkDir, "framework.jar.out", "smali")
        partitionPath = os.path.join(frameworkDir, "framework.jar.out", "smali_classes2")
        if os.path.exists(partitionPath):
            Log.i(TAG, "Combine %s into framework.jar.out/smali" % partitionPath)
            for subDir in os.listdir(partitionPath):
                src = os.path.join(partitionPath, subDir)
                Utils.run(["cp", "-r",  src, dst], stdout=subprocess.PIPE).communicate()

            shutil.rmtree(partitionPath)


        # For Android Handle framework partitions
        for partition in PARTITIONS:
            if partition == "framework.jar.out": continue

            partitionPath = os.path.join(frameworkDir, partition)
            if os.path.exists(partitionPath):
                Log.i(TAG, "Combine %s into framework.jar.out" % partition)
                src = os.path.join(partitionPath, "smali")
                dst = os.path.join(frameworkDir, "framework.jar.out")
                Utils.run(["cp", "-r",  src, dst], stdout=subprocess.PIPE).communicate()
                shutil.rmtree(partitionPath)


    @staticmethod
    def removeUseless(frameworkDir):
        """ Remove useless directory
        """

        for uselessDir in USELESS_DIRS:
            dst = os.path.join(frameworkDir, uselessDir)
            if os.path.exists(dst):
                shutil.rmtree(dst)


    @staticmethod
    def isInFramework(filepath):
        """ Is the file path in jars defined in FRAMEWORK_JARS
        """

        relRoot = filepath.split("/")[0]
        result = relRoot[0:-4] in FRAMEWORK_JARS
        result |= (relRoot == "framework-res")
        Log.d(TAG, "Utils.isInFramework(): %s, %s -> %s" %(result, filepath, relRoot))
        return result


    @staticmethod
    def copyWholly(srcFilePath, dstDirname):
        """ Copy whole SMALI files which are in the same JAVA file
            Especially for the case of inner class, it has '$' in file path
        """

        if not os.path.exists(dstDirname): os.makedirs(dstDirname)

        # Copy all the sub classes even if no change.
        # We need to format the SMALI file
        pos = srcFilePath.find("$")
        if pos > 0:
            srcFilePath = srcFilePath[0:pos] + "*"
        elif srcFilePath.endswith(".smali"):
            srcFilePath = srcFilePath.rstrip(".smali") + "*"

        # Note: Do not use commands.mkarg here
        cmd = "cp %s %s" %(srcFilePath, dstDirname)
        Log.d(TAG, "Utils.copyWholly(): %s" %cmd)
        commands.getstatusoutput(cmd)



    @staticmethod
    def run(args, **kwargs):
        """Create and return a subprocess.Popen object, printing the command
           line on the terminal
        """

        return subprocess.Popen(args, **kwargs)


    @staticmethod
    def printSubprocessOut(subp):
        while True:
            buff = subp.stdout.readline().strip('\n')
            if buff == '' and subp.poll() != None:
                break

            Log.i(TAG, buff)

# End of class Utils


if __name__ == "__main__":
    OPTIONS.handle(sys.argv)
    if OPTIONS.prepare: Prepare()

